/** @file
  Brief description of the file's purpose.
  Detailed description of the file's contents and other useful
  information for a person viewing the file for the first time.
  Copyright (C) 2006 - 2023, Kunlun BIOS, Kunlun Technology (Beijing) Co., Ltd.. All
  Rights Reserved.
  You may not reproduce, distribute, publish, display, perform, modify, adapt,
  transmit, broadcast, present, recite, release, license or otherwise exploit
  any part of this publication in any form, by any means, without the prior
  written permission of Kunlun Technology (Beijing) Co., Ltd..
  @par Revision Reference:18
  - PI Version 1.0
  @par Glossary:
  - IETF - Internet Engineering Task Force
  - NASA - National Aeronautics and Space Administration
**/


#include <Uefi.h>
//[yfliu-0001]
#include <Base.h>
#include <Uefi/UefiBaseType.h>
//[yfliu-0001]
#include <Library/MemoryAllocationLib.h>
#include <Library/BaseLib.h>
#include <Library/DebugLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/PrintLib.h>
#include <Protocol/Smbios.h>
#include <IndustryStandard/SmBios.h>
#include <Guid/DebugMask.h>
//[yfliu-0001]
#include <Library/UefiBootServicesTableLib.h>
#include <Library/KlSmbiosLib.h>
//[yfliu-0001]

/**
  This function is used to set optional string by index.

  @param  OptionalStrStart      A pointer to the optional string start.
  @param  Index                 The index.
  @param  String                A pointer to string.

  @retval EFI_SUCCESS
  @retval EFI_INVALID_PARAMETER

**/
EFI_STATUS
SetOptionalStringByIndex (
  IN CHAR8             *OptionalStrStart,
  IN UINT8             Index,
  IN CHAR8             *String
  )
{
  CHAR8      *LocalStringsBuffer;
  UINT8      StringIndex;
  CHAR8      *OldStrPtr;
  CHAR8      *NewStrPtr;
  UINTN      OldStrSize;
  UINTN      NewStrSize;
  UINTN      TotalSize;

  if ((Index == 0) || (OptionalStrStart == NULL)) {
    return EFI_INVALID_PARAMETER;
  }

  if (String == NULL) {
    return EFI_INVALID_PARAMETER;
  }
  //[gliu-0003]
#define MAX_STRING_SIZE  0x500
  //LocalStringsBuffer = AllocateZeroPool (0x500);
  LocalStringsBuffer = AllocateZeroPool (MAX_STRING_SIZE);
  //[gliu-0003]
  TotalSize = 0;
  OldStrPtr = OptionalStrStart;
  NewStrPtr = LocalStringsBuffer;
  NewStrSize = AsciiStrLen (String) + 1;
  StringIndex = 1;
  do {
    OldStrSize = AsciiStrSize (OldStrPtr);
    if ( StringIndex == Index ) {
    //[gliu-0003]
      AsciiStrCpyS (NewStrPtr, NewStrSize<MAX_STRING_SIZE?NewStrSize:MAX_STRING_SIZE, String);
    //[gliu-0003]
      OldStrPtr += OldStrSize;
      NewStrPtr += NewStrSize;
      TotalSize += NewStrSize;
    }
    else {
    //[gliu-0003]
      AsciiStrCpyS (NewStrPtr, OldStrSize<MAX_STRING_SIZE?OldStrSize:MAX_STRING_SIZE, OldStrPtr);
    //[gliu-0003]
      OldStrPtr += OldStrSize;
      NewStrPtr += OldStrSize;
      TotalSize += OldStrSize;
    }
    StringIndex ++;
  } while (*OldStrPtr != '\0');
  TotalSize ++;
  CopyMem ( OptionalStrStart, LocalStringsBuffer , TotalSize);

  if ( LocalStringsBuffer ) FreePool (LocalStringsBuffer);
  return EFI_SUCCESS;
}

/**
  This function is used to get smbios structure size.

  @param  Head      A pointer to EFI smbios table header.

  @retval Size
  @retval EFI_INVALID_PARAMETER

**/
UINTN
GetSmbiosStructureSize (
  IN EFI_SMBIOS_TABLE_HEADER      *Head
  )
{
  UINTN  Size;
  UINTN  FullSize;
  UINTN  StrLen;
  UINTN  MaxLen;
  INT8*  CharInStr;

  FullSize = Head->Length;
  CharInStr = (INT8*)Head + Head->Length;
  Size = FullSize;
  StrLen = 0;
  //
  // look for the two consecutive zeros, check the string limit by the way.
  //
  while (*CharInStr != 0 || *(CharInStr+1) != 0) {
    if (*CharInStr == 0) {
      Size += 1;
      CharInStr++;
    }

    MaxLen = SMBIOS_TABLE_MAX_LENGTH;

    for (StrLen = 0 ; StrLen < MaxLen; StrLen++) {
      if (*(CharInStr+StrLen) == 0) {
        break;
      }
    }

    if (StrLen == MaxLen) {
      return EFI_INVALID_PARAMETER;
    }

    //
    // forward the pointer
    //
    CharInStr += StrLen;
    Size += StrLen;
  }

  //
  // count ending two zeros.
  //
  Size += 2;
  return Size;
}

/**
  This function is used to init optional strings.

  @param  OptionalStrStart      A pointer to EFI smbios table header.
  @param  StringArray           The string array.

  @retval EFI_SUCCESS

**/
EFI_STATUS
EFIAPI
InitOptionalStrings (
  IN CHAR8              *OptionalStrStart,
  IN CHAR8              **StringArray
  )
{
  UINTN          Index;
  //UINTN          StrSize;


  for (Index = 0; StringArray[Index] != NULL; Index++) {
    //StrSize = AsciiStrSize (StringArray[Index]);
    //[gliu-0001]
    //AsciiStrCpyS (OptionalStrStart, sizeof (OptionalStrStart), StringArray[Index]);
    AsciiStrCpyS (OptionalStrStart, SMBIOS_TABLE_MAX_LENGTH, StringArray[Index]);
    //[gliu-0001]
    OptionalStrStart += AsciiStrSize (OptionalStrStart);
  }
  *OptionalStrStart = '\0';

  return EFI_SUCCESS;
}

//[yfliu-0001]
/**
  Update the SMBIOS Type32.

  @param  BOOTSTATUS

**/
EFI_STATUS
UpdateSmbiosType32 (
  IN      UINT8      BOOTSTATUS
  )
{
  EFI_STATUS                        Status;
  SMBIOS_TABLE_TYPE32               *Type32Record;
  EFI_SMBIOS_TABLE_HEADER           *SmbiosRecord;

  EFI_SMBIOS_PROTOCOL               *gSmbiosProtocol;
  EFI_SMBIOS_HANDLE                 gSmbiosHandle;
  EFI_SMBIOS_TYPE                   gSmbiosType;

  gSmbiosProtocol = NULL;
  gSmbiosHandle    = SMBIOS_HANDLE_PI_RESERVED;
  gSmbiosType      = EFI_SMBIOS_TYPE_SYSTEM_BOOT_INFORMATION;
  Status = gBS->LocateProtocol (
                  &gEfiSmbiosProtocolGuid,
                  NULL,
                  (VOID**)&gSmbiosProtocol
                  );
  if (EFI_ERROR (Status)) {
    gSmbiosProtocol = NULL;
    DEBUG((EFI_D_ERROR,"Get SMBIOS Protocol failed!\n"));
    return EFI_D_ERROR;
  }

  Status = gSmbiosProtocol->GetNext (gSmbiosProtocol, &gSmbiosHandle, &gSmbiosType, &SmbiosRecord, NULL);
  if (EFI_ERROR(Status)) {
    DEBUG((EFI_D_ERROR,"Get SMBIOS Type32 failed!\n"));
    return EFI_D_ERROR;
  }

  Status = gSmbiosProtocol ->Remove(gSmbiosProtocol,gSmbiosHandle);
  gSmbiosHandle = SMBIOS_HANDLE_PI_RESERVED;
  Type32Record = AllocateZeroPool (sizeof (SMBIOS_TABLE_TYPE32) + 1 + 1);
  if (NULL == Type32Record) {
    return EFI_OUT_OF_RESOURCES;
  }

  Type32Record->Hdr.Length = sizeof (SMBIOS_TABLE_TYPE32);
  Type32Record->Hdr.Handle = gSmbiosHandle;
  Type32Record->Hdr.Type = EFI_SMBIOS_TYPE_SYSTEM_BOOT_INFORMATION;
  Type32Record->BootStatus = BOOTSTATUS;
  Status = gSmbiosProtocol ->Add (gSmbiosProtocol,NULL,&gSmbiosHandle,(EFI_SMBIOS_TABLE_HEADER*)Type32Record);
  if (EFI_ERROR (Status)) {
      DEBUG ((EFI_D_ERROR, "Add SMBIOS type 32 failed.\n"));
      return EFI_D_ERROR;
    }

  FreePool (Type32Record);

  return EFI_SUCCESS;
}
//[yfliu-0001]

/**
  This function is used for smbios library to create entry.

  @param  SmbiosEntry      A pointer to smbios structure.
  @param  StringArray      The string array.
  @param  Smbios           A pointer to EFI smbios protocol.
  @param  SmbiosHandle     A pointer to EFI smbios handle.

  @retval EFI_OUT_OF_RESOURCES
  @retval Status

**/
EFI_STATUS
EFIAPI
SmbiosLibCreateEntry (
  IN  SMBIOS_STRUCTURE      *SmbiosEntry,
  IN  CHAR8                 **StringArray,
  IN  EFI_SMBIOS_PROTOCOL   *Smbios,
  OUT EFI_SMBIOS_HANDLE     *SmbiosHandle
  )
{
  EFI_STATUS                Status;
  EFI_SMBIOS_TABLE_HEADER   *Record;
  UINTN                     Index;
  UINTN                     StringSize;
  UINTN                     Size;
  CHAR8                     *Str;

  //
  // Calculate the size of the fixed record and optional string pack
  //
  Size = SmbiosEntry->Length;
  if (StringArray == NULL) {
    Size += 2; // Min string section is double null
  } else if (StringArray[0] == NULL) {
    Size += 2; // Min string section is double null
  } else {
    for (Index = 0; StringArray[Index] != NULL; Index++) {
      StringSize = AsciiStrSize (StringArray[Index]);
      Size += StringSize;
    }
    // Don't forget the terminating double null
    Size += 1;
  }

  // Copy over Template
  Record = (EFI_SMBIOS_TABLE_HEADER *)AllocateZeroPool (Size);
  if (Record == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  CopyMem (Record, SmbiosEntry, SmbiosEntry->Length);

  if (StringArray != NULL) {
    // Append string pack
    Str = ((CHAR8 *)Record) + Record->Length;
    for (Index = 0; StringArray[Index] != NULL; Index++) {
      StringSize = AsciiStrSize (StringArray[Index]);
      CopyMem (Str, StringArray[Index], StringSize);
      Str += StringSize;
    }
    *Str = 0;
  }

  *SmbiosHandle = SMBIOS_HANDLE_PI_RESERVED;
  Status = Smbios->Add (
                     Smbios,
                     NULL,
                     SmbiosHandle,
                     Record
                     );
  FreePool (Record);

  return Status;
}
